import os
import json
import yaml
import argparse
import numpy as np
import rospkg


def parse_list(arg):
    """将命令行参数解析为列表"""
    return list(map(int, arg.strip('[]').split(',')))


def read_yaml_file(yaml_file):
    if not os.path.exists(yaml_file):
        return None
    with open(yaml_file, 'r', encoding='utf-8') as file:
        data = yaml.safe_load(file)
    return data


def write_yaml_file(yaml_data, yaml_file):
    # 将Python字典转换为YAML格式字符串
    yaml_str = yaml.dump(yaml_data, default_flow_style=False, sort_keys=False, allow_unicode=True, indent=4)
    # 将YAML数据写入文件
    with open(yaml_file, 'w', encoding='utf-8') as file:
        file.write(yaml_str)


def get_ros_package_path(package_name):
    try:
        ros_pack = rospkg.RosPack()
        return ros_pack.get_path(package_name)
    except rospkg.common.ResourceNotFound:
        return None


def read_calib_param(calib_file):
    if not os.path.exists(calib_file):
        raise ValueError("calib_path does not exist")
    with open(calib_file, 'r', encoding='utf-8') as file:
        # 将JSON字符串解析为Python字典
        json_data = json.load(file)

    camera_list = []
    lidar_dict = {}
    for i in range(len(json_data["cameras"])):
        camera_param = json_data["cameras"][i]
        camera_list.append({
            "name": camera_param['name'],
            "T_lidar_cam": [
                camera_param["transform_from_lidar"]["rotation"][0][0],
                camera_param["transform_from_lidar"]["rotation"][0][1],
                camera_param["transform_from_lidar"]["rotation"][0][2],
                camera_param["transform_from_lidar"]["position"][0],
                camera_param["transform_from_lidar"]["rotation"][1][0],
                camera_param["transform_from_lidar"]["rotation"][1][1],
                camera_param["transform_from_lidar"]["rotation"][1][2],
                camera_param["transform_from_lidar"]["position"][1],
                camera_param["transform_from_lidar"]["rotation"][2][0],
                camera_param["transform_from_lidar"]["rotation"][2][1],
                camera_param["transform_from_lidar"]["rotation"][2][2],
                camera_param["transform_from_lidar"]["position"][2],
                0.0, 0.0, 0.0, 1.0
            ],
            "distortion_coeffs": list(camera_param["distortion"]["params"].values()),
            "distortion_model": "equidistant",
            "intrinsics": [
                camera_param["intrinsic"]["fl_x"],
                camera_param["intrinsic"]["fl_y"],
                camera_param["intrinsic"]["cx"],
                camera_param["intrinsic"]["cy"]
            ],
            "resolution": [camera_param["width"], camera_param["height"]],
        })

    if "imu" in json_data:
        lidar2imu = json_data["imu"]["transform_from_lidar"]
        lidar_dict = {
            "T_imu_lidar": [
                lidar2imu['rotation'][0][0], lidar2imu['rotation'][0][1], lidar2imu['rotation'][0][2], lidar2imu["position"][0],
                lidar2imu['rotation'][1][0], lidar2imu['rotation'][1][1], lidar2imu['rotation'][1][2], lidar2imu["position"][1],
                lidar2imu['rotation'][2][0], lidar2imu['rotation'][2][1], lidar2imu['rotation'][2][2], lidar2imu["position"][2],
                0.0, 0.0, 0.0, 1.0
            ],
        }
    else:
        lidar_dict = {
            "T_imu_lidar": [
                1.0, 0.0, 0.0, 0.011,
                0.0, 1.0, 0.0, 0.02329,
                0.0, 0.0, 1.0, -0.04412,
                0.0, 0.0, 0.0, 1.0
            ],
        }
    calib_param = {}
    calib_param["camera"] = camera_list
    calib_param["lidar"] = lidar_dict
    return calib_param


def resize_camera_param(camera_param, resolution):
    kx = float(resolution[0]) / float(camera_param['resolution'][0])
    ky = float(resolution[1]) / float(camera_param['resolution'][1])
    camera_param['intrinsics'][0] *= kx
    camera_param['intrinsics'][1] *= ky
    camera_param['intrinsics'][2] *= kx
    camera_param['intrinsics'][3] *= ky
    camera_param['resolution'] = resolution
    return camera_param


def change_adapter_param(calib_param: dict, adapter_param: dict, out_resolution: list):
    if adapter_param is None:
        raise ValueError("adapter_param is None, please check the YAML file path or content.")

    # 重新组织数据以匹配目标YAML格式
    camera_dict = {}
    camera_params = calib_param["camera"]
    adapter_param["max_cameras"] = len(camera_params)
    for i in range(len(camera_params)):
        camera_param = camera_params[i]
        camera_dict["cam" + str(i)] = {
            "T_lidar_cam": camera_param["T_lidar_cam"],
            "intrinsics": camera_param["intrinsics"],
            "resolution": camera_param["resolution"],
            "distortion_coeffs": camera_param["distortion_coeffs"],
            "cam_overlaps": [1-i],
            "camera_model": "pinhole",
            "distortion_model": "equidistant",
            "rostopic": f"/camera/{camera_param['name']}/jpeg",
            "rostopic_out": f"/camera/{camera_param['name']}/jpeg_1k/data",
            "rostopic_out_undistort": f"/camera/{camera_param['name']}/jpeg_1k/undistort",
            "out_resolution": out_resolution,
            "timeshift_cam_imu": 0.0
        }

    adapter_param.update(camera_dict)

    adapter_param["lidar0"] = {
        "rostopic": "/livox/lidar"
    }
    adapter_param["lidar0"].update(calib_param["lidar"])


def change_fastlivo_lidar_param(lidar_param, adapter_param):
    T_imu2lidar = np.array(adapter_param["lidar0"]["T_imu_lidar"]).reshape(4, 4)
    T_lidar2imu = np.linalg.inv(T_imu2lidar)
    lidar_param['mapping']['extrinsic_T'] = T_lidar2imu[:3, 3].ravel().tolist()
    lidar_param['mapping']['extrinsic_R'] = T_lidar2imu[:3, :3].ravel().tolist()
    T_lidar2cam = np.array(adapter_param["cam0"]["T_lidar_cam"]).reshape(4, 4)
    lidar_param["camera"]["Pcl"] = T_lidar2cam[:3, 3].ravel().tolist()
    lidar_param["camera"]["Rcl"] = T_lidar2cam[:3, :3].ravel().tolist()


def generate_fastlivo_intrinsic(adapter_param):
    resized_intrinsic_param = {}
    for key in adapter_param.keys():
        if isinstance(key, str) and key.startswith("cam"):
            camera_param = adapter_param[key]
            kx = float(camera_param['out_resolution'][0]) / float(camera_param['resolution'][0])
            ky = float(camera_param['out_resolution'][1]) / float(camera_param['resolution'][1])
            resized_intrinsic_param[key] = {
                'cam_fx': camera_param['intrinsics'][0] * kx,
                'cam_fy': camera_param['intrinsics'][1] * ky,
                'cam_cx': camera_param['intrinsics'][2] * kx,
                'cam_cy': camera_param['intrinsics'][3] * ky,
                'cam_width': camera_param['out_resolution'][0],
                'cam_height': camera_param['out_resolution'][1],
                'cam_d0': 0.0,
                'cam_d1': 0.0,
                'cam_d2': 0.0,
                'cam_d3': 0.0,
            }
    return resized_intrinsic_param


def generate_filter_param(camera_params):
    filter_param = {}
    for key, value in (camera_params.items()):
        if not key.startswith("cam"):
            raise ValueError("输入的字符串必须以 'cam' 开头")

        num_part = key[3:]
        new_num = int(num_part) + 1

        filter_param[f"cam_{new_num}"] = {
            'image_width': value['cam_width'],
            'image_height': value['cam_height'],
            'camera_intrinsic': [
                value['cam_fx'], 0.0, value['cam_cx'],
                0.0, value['cam_fy'], value['cam_cy'],
                0.0, 0.0, 1.0
            ],
            'camera_dist_coeffs': [0.0, 0.0, 0.0, 0.0, 0.0],
            'time_diff_camera_to_lidar': 0.0
        }
    return filter_param


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="convert parameters from json to yaml")
    adapter_dir = get_ros_package_path('adapter')
    fastlivo_dir = get_ros_package_path('fast_livo')
    fastlio_dir = get_ros_package_path('fast_lio')
    data_preprocessing_dir = get_ros_package_path('data_preprocessing')

    parser.add_argument('--calib-path', type=str, default='', help="The calibration result of device")
    parser.add_argument('--in-resolution', type=parse_list, default=[3040, 4032], help="Input resolution in the format [width, height]")
    parser.add_argument('--out-resolution', type=parse_list, default=[760, 1008], help="Output resolution in the format [width, height]")
    args = parser.parse_args()
    calib_param = read_calib_param(args.calib_path)

    for i in range(len(calib_param["camera"])):
        calib_param["camera"][i] = resize_camera_param(calib_param["camera"][i], args.in_resolution)

    if adapter_dir is None:
        raise ValueError("adapter package not found")

    adapter_param_path = os.path.join(adapter_dir, "config/fast_livo.yaml")
    if not os.path.exists(adapter_param_path):
        raise ValueError("adapter_param_path does not exist")
    adapter_param = read_yaml_file(adapter_param_path)

    change_adapter_param(calib_param, adapter_param, args.out_resolution)
    write_yaml_file(adapter_param, adapter_param_path)

    if data_preprocessing_dir is not None:
        write_yaml_file(adapter_param, os.path.join(data_preprocessing_dir, "config/data_preprocessing.yaml"))

    resized_cameras_param = generate_fastlivo_intrinsic(adapter_param)

    if fastlivo_dir is not None:
        livo_camera_yaml = os.path.join(fastlivo_dir, "config/metacam_camera.yaml")
        livo_camera_param = read_yaml_file(livo_camera_yaml)
        livo_camera_param.update(resized_cameras_param)
        write_yaml_file(livo_camera_param, livo_camera_yaml)

        livo_lidar_yaml = os.path.join(fastlivo_dir, "config/metacam_lidar.yaml")
        livo_lidar_param = read_yaml_file(livo_lidar_yaml)
        change_fastlivo_lidar_param(livo_lidar_param, adapter_param)
        write_yaml_file(livo_lidar_param, livo_lidar_yaml)
    filter_param_path = os.path.join(os.path.dirname(os.path.dirname(args.calib_path)), "metacam_intrinsic_extrinsic.yaml")
    write_yaml_file(generate_filter_param(resized_cameras_param), filter_param_path)